# @Time    : 2018/12/31 20:40
# @Author  : heyin
import pandas as pd
from sklearn.linear_model import LinearRegression, SGDRegressor
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler


def line():
    """线性回归"""
    df = pd.read_csv('./data.csv')
    df = df[df['vote_count'] > 2000]
    df.dropna(axis=0, how='any', inplace=True)
    y = df[['score']]  # series，stars作为目标值
    # df = df.loc[:, 'year':]
    x = df.drop('stars', axis=1)  # 删掉stars后其余的作为特征值
    for i in range(1, 6):
        x['actor%s' % i] = x['actor%s_best' % i] * 0.4 + x['actor%s_worst' % i] * 0.2 + x['actor%s_recent' % i] * 0.3
        x = x.drop(['actor%s_best' % i, 'actor%s_worst' % i, 'actor%s_recent' % i], axis=1)
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.25, random_state=1)
    print(x_train.columns)
    x_train_title = x_train['title']
    x_test_title = x_test['title']
    x_train = x_train.loc[:, 'year':]
    x_test = x_test.loc[:, 'year':]
    # print(x_test_title)

    sd_x = StandardScaler()
    x_train = sd_x.fit_transform(x_train)
    x_test = sd_x.transform(x_test)

    sd_y = StandardScaler()
    y_train = sd_y.fit_transform(y_train)

    # 进行估计器的运算
    # 正规方程
    lr = LinearRegression()
    lr.fit(x_train, y_train)
    # 预测房子价格要的是真实价格，因此需要将数据逆转
    y_prelr = lr.predict(x_test)
    pred_score = sd_y.inverse_transform(y_prelr)

    y_pre_train = lr.predict(x_train)
    y_pre_train = sd_y.inverse_transform(y_pre_train)

    pred_score = pd.DataFrame(pred_score, columns=['pred_score'], index=x_test_title.index)
    test_score = pd.concat([x_test_title, pred_score], axis=1)
    print(test_score)
    train_score = pd.DataFrame(y_pre_train, columns=['pred_score'], index=x_train_title.index)
    train_score = pd.concat([x_train_title, train_score], axis=1)
    print(train_score)

    ret = pd.concat([test_score, train_score], axis=0)
    ret.to_csv('./ret.csv', index=False)

    # 均方误差评判效果
    mse1 = mean_squared_error(y_test, sd_y.inverse_transform(y_prelr))
    mse2 = mean_squared_error(sd_y.inverse_transform(y_train), y_pre_train)
    print(mse1, mse2)


def sgd():
    df = pd.read_csv('./data.csv')
    df = df[df['vote_count'] > 2000]
    df.dropna(axis=0, how='any', inplace=True)
    y = df[['score']]  # series，stars作为目标值
    # df = df.loc[:, 'year':]
    x = df.drop('stars', axis=1)  # 删掉stars后其余的作为特征值
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.25, random_state=1)
    # print(x_train)
    x_train_title = x_train['title']
    x_test_title = x_test['title']
    x_train = x_train.loc[:, 'year':'actor5_recent']
    x_test = x_test.loc[:, 'year':'actor5_recent']
    # print(x_test_title)

    sd_x = StandardScaler()
    x_train = sd_x.fit_transform(x_train)
    x_test = sd_x.transform(x_test)

    sd_y = StandardScaler()
    y_train = sd_y.fit_transform(y_train)

    sgdr = SGDRegressor(max_iter=5)
    sgdr.fit(x_train, y_train)

    y_prelr = sgdr.predict(x_test)
    pred_score = sd_y.inverse_transform(y_prelr)

    y_pre_train = sgdr.predict(x_train)
    y_pre_train = sd_y.inverse_transform(y_pre_train)

    # pred_score = pd.DataFrame(pred_score, columns=['pred_score'], index=x_test_title.index)
    # test_score = pd.concat([x_test_title, pred_score], axis=1)
    # print(test_score)
    # train_score = pd.DataFrame(y_pre_train, columns=['pred_score'], index=x_train_title.index)
    # train_score = pd.concat([x_train_title, train_score], axis=1)
    # print(train_score)

    # ret = pd.concat([test_score, train_score], axis=0)
    # ret.to_csv('./ret.csv', index=False)

    # # 均方误差评判效果
    mse = mean_squared_error(y_test, sd_y.inverse_transform(y_prelr))
    print(mse)


if __name__ == '__main__':
    line()
    # sgd()
